package com.txy.redblack;

/**
 * 每个节点要么是黑色，要么是红色
 * 根节点是黑色
 * 每个叶子节点是黑色的空节点，简单起见有时会省略
 * 如果一个节点是红色的，则它的子节点必须是黑色，父子不能都为红
 * 从任一节点到其每个叶子的所有路径都包含相同数目的黑色节点
 * 新插入节点默认为红色，插入后校验红黑树是否符合规则，不符合的要进行平衡
 */
public class RBTree {
    RBTreeNode root;

    public void list(RBTreeNode node) {
        if(node == null) {
            return;
        }
        if(node.getLeft() == null && node.getRight() == null) {
            System.out.println(node);
            return ;
        }
        System.out.println(node);
        list(node.getLeft());
        list(node.getRight());
    }

    public void insert(int key) {
        RBTreeNode node = new RBTreeNode(key);
        if(root==null) {
            node.setBlack();
            root = node;
            return;
        }
        RBTreeNode parent = root;
        RBTreeNode son = null;
        if(key <= parent.getKey()) {
            son = parent.getLeft();
        }else {
            son = parent.getRight();
        }
        while(son != null) {
            parent = son;
            if(key <= parent.getKey()) {
                son = parent.getLeft();
            } else {
                son = parent.getRight();
            }
        }
        if(key <= parent.getKey()) {
            parent.setLeft(node);
        }else {
            parent.setRight(node);
        }
        node.setParent(parent);

        // 自平衡
        balanceInsert(node);
    }

    private void balanceInsert(RBTreeNode node) {
        RBTreeNode father, grandFather;
        while((father = node.getParent())!=null && father.isBlack()==false) {
            grandFather = father.getParent();
            if(grandFather.getLeft() == father) {
                RBTreeNode uncle = grandFather.getRight();
                if(uncle!=null && uncle.isRed()) {
                    father.setBlack();
                    uncle.setBlack();
                    grandFather.setRed();
                    node = grandFather;
                    continue;
                }
                if(node==father.getRight()) {
                    leftRotate(father);
                    RBTreeNode tmp = node;
                    node = father;
                    father = tmp;
                }
                father.setBlack();
                grandFather.setRed();
                rightRotate(grandFather);
            }else {
                RBTreeNode uncle = grandFather.getLeft();
                if(uncle!=null && uncle.isRed()) {
                    father.setBlack();
                    uncle.setBlack();
                    grandFather.setRed();
                    node = grandFather;
                    continue;
                }
                if(node == father.getLeft()) {
                    rightRotate(father);
                    RBTreeNode tmp = node;
                    node = father;
                    father = tmp;
                }
                father.setBlack();
                grandFather.setRed();
                leftRotate(grandFather);
            }
        }
        root.setBlack();
    }

    private void leftRotate(RBTreeNode node) {
        RBTreeNode right = node.getRight();
        RBTreeNode parent = node.getParent();
        if(parent == null) {
            root = right;
            right.setParent(null);
        }else {
            if(parent.getLeft()!=null && parent.getLeft()==node) {
                parent.setLeft(right);
            }else {
                parent.setRight(right);
            }
            right.setParent(parent);
        }
        node.setParent(right);
        node.setRight(right.getLeft());
        if(right.getLeft()!=null) {
            right.getLeft().setParent(node);
        }
        right.setLeft(node);
    }

    private void rightRotate(RBTreeNode node) {
        RBTreeNode left = node.getRight();
        RBTreeNode parent = node.getParent();
        if(parent == null) {
            root = left;
            left.setParent(null);
        } else {
            if(parent.getLeft() != null && parent.getLeft() == node) {
                parent.setLeft(left);
            } else {
                parent.setRight(left);
            }
            left.setParent(parent);
        }
        node.setParent(left);
        node.setLeft(left.getRight());
        if(left.getRight()!=null) {
            left.getRight().setParent(node);
        }
        left.setRight(node);
    }
}

